package medium;

import java.util.ArrayList;
import java.util.List;

public class Lc2698 {
    class Solution {
        private List<Integer> splitPos = new ArrayList<>();
        public int punishmentNumber(int n) {
            int ret = 0;
            for (int i = 1; i <= n; i++) {
                if (this.isGoodNum(i)) {
                    ret +=  i * i;
                }
            }
            return ret;
        }

        private boolean isGoodNum(int n) {
            this.splitPos.clear();
            this.splitPos.add(-1);
            return this.dfs(String.valueOf(n * n), n, 0);
        }

        // 对ori字符串的所有可切分位置进行回溯
        // 切分位置指对idx位置字符后面的空位
        // 切分后得到的字符串位ori[0, pos + 1]
        // 结束条件为pos == ori.len，注意取字符串时pos要+1
        // index指示当前位置，范围0 - ori.length - 1，同pos
        private boolean dfs(String ori, int target, int index) {
            // 边界条件
            // 最后一个位置，一定取，取后计算结果
            if (index == ori.length() - 1) {
                return this.check(ori, target);
            }

            // 取
            this.splitPos.add(index);
            // 继续寻找下一个位置
            boolean ret1 = dfs(ori, target, index + 1);
            // 回退
            this.splitPos.remove(this.splitPos.size() - 1);

            // 不取
            boolean ret2 = dfs(ori, target, index + 1);
            return ret1 || ret2;
        }

        private boolean check(String ori, int target) {
            int sum = 0;
            for (int i = 1; i < this.splitPos.size(); i++) {
                sum += Integer.parseInt(ori.substring(this.splitPos.get(i - 1) + 1, this.splitPos.get(i) + 1));
            }
            sum += Integer.parseInt(ori.substring(this.splitPos.get(this.splitPos.size() - 1) + 1));
            return sum == target;
        }
    }
}
